import json
import os
import time

import javalang
from app.logic.utils.filepath import TOKEN_DIR, SRC_DIR

'''
按项目路径生成token序列
'''


def extract_tokens(java_file_path):
    """
    从 Java 文件中提取 token 令牌序列
    """
    try:
        with open(java_file_path, 'r', encoding='utf-8') as f:
            code = f.read()
        # 使用 javalang 解析 Java 代码
        tokens = []
        # 使用 javalang.tokenizer.Tokenizer 来生成 token 序列
        for token in javalang.tokenizer.tokenize(code):
            tokens.append(token.value)  # 获取 token 的值
        return tokens
    except Exception as e:
        print(f"Failed to extract tokens from {java_file_path}: {e}")
        return []


class GenerateToken:
    def __init__(self, work_path_dir, project_path_dir):
        self.workspace = os.path.join(work_path_dir, str(project_path_dir))
        self.src_dir = os.path.join(work_path_dir, str(project_path_dir), SRC_DIR)
        self.token_dir = os.path.join(work_path_dir, str(project_path_dir), TOKEN_DIR)

        # 确保 token 存储目录存在
        if not os.path.exists(self.token_dir):
            os.makedirs(self.token_dir)

    def save_tokens_for_file(self, java_file_path):
        """
        生成单个文件的 token 令牌序列并保存为 JSON 文件
        """
        tokens = extract_tokens(java_file_path)
        if tokens:
            # 获取文件名并生成保存路径
            file_name = os.path.basename(java_file_path)  # 获取文件名
            token_file = os.path.join(self.token_dir, f"{file_name}.json")
            # 将 token 列表保存为 JSON 文件
            with open(token_file, 'w', encoding='utf-8') as f:
                json.dump(tokens, f, indent=4, ensure_ascii=False)
            print(f"Tokens saved to {token_file}.")

    def generate_tokens_for_directory(self):
        """
        遍历目录中的所有 Java 文件并生成它们的 token 令牌序列
        """
        for root, _, files in os.walk(self.src_dir):
            for file in files:
                if file.endswith('.java'):
                    java_file = os.path.join(root, file)
                    self.save_tokens_for_file(java_file)

    def extract_and_save_tokens(self):
        """
        打印开始时间，生成 token 令牌序列并保存，最后打印结束时间和总耗时
        """
        start_time = time.time()
        print("Token Extraction Start Time: " + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        # 生成 token 令牌序列
        self.generate_tokens_for_directory()

        # 打印结束时间
        print("Token extraction completed at " + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

        # 计算总耗时
        total_time = time.time() - start_time
        print("Total time: " + str(total_time // 3600) + "h" + str((total_time % 3600) // 60) + "m" + str(
            (total_time % 3600) % 60) + "s")


def gen_tokens(work_dir, project_path):
    token_generator = GenerateToken(work_dir, project_path)
    token_generator.generate_tokens_for_directory()


# 使用方法示例
if __name__ == "__main__":
    work_dir = '../../../data'  # 请修改为实际路径
    project_path = '29'  # 请修改为实际项目路径

    token_generator = GenerateToken(work_dir, project_path)
    token_generator.extract_and_save_tokens()
